package edu.buaa.common.client;

import com.google.common.util.concurrent.ListenableFuture;
import edu.buaa.common.transaction.AbstractTransaction;
import edu.buaa.common.transaction.TransactionFailedException;
import edu.buaa.utils.TimeMonitor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * Created by crusher. 2020.10.14
 */
public abstract class AbstractSQLClient implements DBProxy {
    private static final Logger log = LoggerFactory.getLogger(AbstractSQLClient.class);
    protected final String dbName;
    protected final String serverHost;
    protected final BlockingQueue<Connection> connectionPool = new LinkedBlockingQueue<>();
    protected final Map<Connection, Integer> connIdMap = new HashMap<>();

    private final RequestDispatcher service;

    public AbstractSQLClient(String serverHost, String dbName, int parallelCnt) throws Exception {
        this.serverHost = serverHost;
        this.dbName = dbName;
        this.service = new RequestDispatcher(parallelCnt);
        createDbIfNotExist(serverHost, dbName);
        for (int i = 0; i < parallelCnt; i++) {
            Connection conn = createNormalConnection(serverHost, dbName);
            this.connectionPool.offer(conn);
            connIdMap.put(conn, i);
            if(i==parallelCnt-1) connected(conn);
        }
    }

    protected abstract void connected(Connection conn) throws Exception;

    protected abstract void createDbIfNotExist(String serverHost, String dbName) throws SQLException, ClassNotFoundException;

    protected abstract List<String> createTables();

    protected abstract Connection createNormalConnection(String serverHost, String dbName) throws SQLException;

    public abstract String currentStorageStatus() throws SQLException, InterruptedException;

//    public abstract void onError(AbstractTransaction tx, Exception e);

    @Override
    public String testServerClientCompatibility() throws UnsupportedOperationException {
        return null;
    }

    @Override
    public void close() throws IOException, InterruptedException {
        try {
            int remains = service.getQueueSize();
            log.debug("Client closed but will send remaining {} requests", remains);
            long mark = System.currentTimeMillis();
            while (remains>0) {
                long completeCnt = service.getCompletedTaskCount();
                remains = service.getQueueSize();
                System.out.println(completeCnt + " / " + (completeCnt + remains) + " query completed.");
                Thread.sleep(20_000);
            }
            service.awaitClose();
            while (!connectionPool.isEmpty()) {
                Connection conn = connectionPool.take();
                conn.close();
            }
            log.debug("Client exit after wait {} seconds. send {} lines.", (System.currentTimeMillis()-mark)/1000, service.getCompletedTaskCount());
        } catch (SQLException e) {
            e.printStackTrace();
            throw new IOException(e);
        }
    }

    protected ListenableFuture<ServerResponse> submit(Req request) throws InterruptedException {
        return this.service.submit(new SQLReq(connectionPool, connIdMap, request), -1);
    }

    protected ListenableFuture<ServerResponse> submit(Req request, int section) throws InterruptedException {
        return this.service.submit(new SQLReq(connectionPool, connIdMap, request), section);
    }

    @FunctionalInterface
    protected interface Req{
        AbstractTransaction.Result executeQuery(Connection conn) throws Exception;
    }

    protected static class SQLReq implements RequestDispatcher.RequestWs {
        private final TimeMonitor timeMonitor = new TimeMonitor();
        protected final AbstractTransaction.Metrics metrics = new AbstractTransaction.Metrics();
        protected final BlockingQueue<Connection> connectionPool;
        protected final Map<Connection, Integer> connIdMap;
        private final Req body;

        protected SQLReq(BlockingQueue<Connection> connectionPool, Map<Connection, Integer> connIdMap, Req body) {
            this.connectionPool = connectionPool;
            this.connIdMap = connIdMap;
            this.body = body;
            timeMonitor.begin("Wait in queue");
        }

        @Override
        public ServerResponse call() throws Exception {
            try {
                Connection conn = connectionPool.take();
                timeMonitor.mark("Wait in queue", "query");
                AbstractTransaction.Result result = body.executeQuery(conn);
                timeMonitor.end("query");
                if (result == null) throw new RuntimeException("[Got null. Server close connection]");
                connectionPool.put(conn);
                metrics.setWaitTime(Math.toIntExact(timeMonitor.duration("Wait in queue")));
                metrics.setSendTime(timeMonitor.beginT("query"));
                metrics.setExeTime(Math.toIntExact(timeMonitor.duration("query")));
                metrics.setConnId(connIdMap.get(conn));
                metrics.setTxSuccess(true);
                ServerResponse response = new ServerResponse();
                response.setMetrics(metrics);
                response.setResult(result);
                return response;
            }catch (TransactionFailedException e){
                ServerResponse response = new ServerResponse();
                metrics.setTxSuccess(false);
                response.setMetrics(metrics);
                response.setResult(null);
                return response;
            } catch (Exception e) {
                e.printStackTrace();
                throw e;
            }
        }
    }
}

